from ultralytics import YOLO
import cv2
import os
from pathlib import Path


def detect_and_crop_images(model_path, input_folder, output_folder, target_classes=[0, 1]):
    """
    使用YOLOv8模型检测图片并裁剪包含目标类别的图片

    参数:
        model_path: 训练好的YOLOv8模型路径
        input_folder: 待检测图片文件夹
        output_folder: 输出文件夹(保存裁剪的图片)
        target_classes: 需要检测的目标类别列表(默认为[0, 1])
    """
    # 加载模型
    model = YOLO(model_path)

    # 确保输出文件夹存在
    os.makedirs(output_folder, exist_ok=True)

    # 获取输入文件夹中的所有图片文件
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
    image_files = [f for f in os.listdir(input_folder)
                   if os.path.splitext(f)[1].lower() in image_extensions]

    # 处理每张图片
    for image_file in image_files:
        image_path = os.path.join(input_folder, image_file)

        # 使用模型进行预测
        results = model(image_path)

        # 遍历检测结果
        for result in results:
            # 获取检测到的类别和边界框
            boxes = result.boxes
            for box in boxes:
                class_id = int(box.cls)
                if class_id in target_classes:
                    # 读取原始图片
                    img = cv2.imread(image_path)

                    # 获取边界框坐标
                    x1, y1, x2, y2 = map(int, box.xyxy[0])

                    # 裁剪图片
                    cropped_img = img[y1:y2, x1:x2]

                    # 生成输出文件名
                    base_name = os.path.splitext(image_file)[0]
                    output_name = f"{base_name}_class{class_id}_{len(os.listdir(output_folder))}.jpg"
                    output_path = os.path.join(output_folder, output_name)

                    # 保存裁剪后的图片
                    cv2.imwrite(output_path, cropped_img)
                    print(f"Saved cropped image: {output_path}")


if __name__ == "__main__":
    # 配置参数
    model_path = "E:/CodeCNN/yolov8-study/runs/detect/train29/weights/best.pt"  # 替换为你的模型路径
    input_folder = r"C:\Users\Administrator\Desktop\testPic"  # 替换为输入图片文件夹
    output_folder = r"C:\Users\Administrator\Desktop\testPic\cropped"  # 替换为输出文件夹

    # 执行检测和裁剪
    detect_and_crop_images(model_path, input_folder, output_folder)